import math
from scipy.interpolate import CubicSpline
import numpy as np


class Spline:
    def __init__(self, distancePreStep: float) -> None:
        self.theta = []
        self.v = []
        self.x = []
        self.y = []
        self.s = []
        self.length = 0
        self.distancePreStep = distancePreStep

    def drawPath(self, ax, color: str, index: bool = False, scater: bool = False):
        import matplotlib.pyplot as plt

        ax.plot(self.x, self.y, color=color)
        if scater:
            ax.scatter(self.x, self.y, color=color)
        if index:
            for i in range(len(self.x)):
                plt.text(self.x[i], self.y[i], str(i))


class ConstVelocitySpline(Spline):
    def __init__(
        self, points: list[tuple[float, float]], velocity: float, distancePreStep: float
    ) -> None:
        """
        points: list[tuple[float,float]], [[x1,y1],[x2,y2],...]
        """
        super().__init__(distancePreStep)
        dis = [0.0]
        for i in range(1, len(points)):
            dis.append(
                dis[i - 1]
                + math.sqrt(
                    (points[i][0] - points[i - 1][0]) ** 2
                    + (points[i][1] - points[i - 1][1]) ** 2
                )
            )
        cs = CubicSpline(dis, points)
        self.length = dis[len(dis) - 1]
        samplerS = np.linspace(0, self.length, int(self.length / self.distancePreStep))
        points = [(p[0], p[1]) for p in cs(samplerS)]

        for i in range(len(points)):
            p = points[i]
            self.x.append(p[0])
            self.y.append(p[1])
            self.v.append(velocity)
            if i > 0:
                dx = points[i][0] - points[i - 1][0]
                dy = points[i][1] - points[i - 1][1]
                self.theta.append(math.atan2(dy, dx))
                self.s.append(self.s[i - 1] + math.sqrt(dx**2 + dy**2))
            else:
                dx = points[1][0] - points[0][0]
                dy = points[1][1] - points[0][1]
                self.theta.append(math.atan2(dy, dx))
                self.s.append(0)


if __name__ == "__main__":
    points: list[tuple[float, float]] = [(0, 0), (1, 2), (2, 0), (3, 4), (10, 0)]
    spline = ConstVelocitySpline(points, 1, 0.5)

    from matplotlib import pyplot as plt

    ax = plt.subplot(1, 1, 1)
    spline.drawPath(ax, "blue", True)
    plt.show()
